# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: neo, jaewon

import datetime
import pytz
from typing import List, Optional, Tuple, Union


class DateFormatError(Exception):
  pass


def to_datetime(timestamp: Union[datetime.datetime, datetime.date, int]) \
    -> datetime.datetime:
  if isinstance(timestamp, int):
    return datetime.datetime.utcfromtimestamp(timestamp / 10.**9)
  elif isinstance(timestamp, datetime.datetime):
    return timestamp
  elif isinstance(timestamp, datetime.date):
    return datetime.datetime(timestamp.year, timestamp.month, timestamp.day)
  else:
    raise ValueError()


def to_timestamp_int(timestamp: Union[int, datetime.datetime, datetime.date]):
  if isinstance(timestamp, int):
    return timestamp
  elif isinstance(timestamp, datetime.datetime):
    return nanoseconds_from_epoch(timestamp)
  elif isinstance(timestamp, datetime.date):
    timestamp = datetime.datetime(timestamp.year, timestamp.month, timestamp.day)
    return nanoseconds_from_epoch(timestamp)
  else:
    raise ValueError()


def nanoseconds_from_epoch(datetime_in: datetime.datetime) -> int:
  if not isinstance(datetime_in, datetime.datetime):
    raise TypeError()
  td = datetime_in - datetime.datetime.fromtimestamp(0, datetime_in.tzinfo)
  seconds = td.seconds + td.days * 24 * 3600
  nanoseconds = td.microseconds * 10**3 + seconds * 10**9
  # 0x7fffffffffffffff is year 2262.
  if not 0x7fffffffffffffff >= nanoseconds >= -0x8000000000000000:
    raise OverflowError('0x%x is out of range (min_int64, max_int64)' % nanoseconds)
  return nanoseconds


def parse_end_time(end_time_string=None, now=None, margin=None):
  if end_time_string is None:
    return None
  splitted = end_time_string.split(':')
  assert len(splitted) == 3 or len(splitted) == 2
  now = now or datetime.datetime.now()
  margin = margin or datetime.timedelta(minutes=10)
  end_time = now.replace(hour=int(splitted[0]),
                         minute=int(splitted[1]),
                         second=int(0 if len(splitted) == 2 else splitted[2]),
                         microsecond=0)
  if now + margin >= end_time:
    end_time += datetime.timedelta(days=1)
    assert now < end_time
  return end_time


def convert_string_to_dates(date_specifier):
  if not isinstance(date_specifier, str):
    raise DateFormatError('Invalid data type: %s' % type(date_specifier))

  res = []
  date_chunks = date_specifier.split(',')
  for date_chunk in date_chunks:
    date_pair = date_chunk.strip().split('-')
    if len(date_pair) == 1:
      if len(date_pair[0]) != 8:
        raise DateFormatError("Wrong date format: %s" % date_chunk)
      year = int(date_pair[0][:4])
      month = int(date_pair[0][4:6])
      day = int(date_pair[0][6:8])
      begin_date = datetime.date(year, month, day)
      end_date = datetime.date(year, month, day)
    elif len(date_pair) == 2:
      if len(date_pair[0]) != 8 or len(date_pair[1]) != 8:
        raise DateFormatError("Wrong date format: %s" % date_chunk)
      begin_date = datetime.date(int(date_pair[0][:4]),
                                 int(date_pair[0][4:6]),
                                 int(date_pair[0][6:8]))
      end_date = datetime.date(int(date_pair[1][:4]),
                               int(date_pair[1][4:6]),
                               int(date_pair[1][6:8]))
    else:
      raise DateFormatError("Wrong date format: %s" % date_chunk)
    add_delta = datetime.timedelta(days=1)
    while begin_date <= end_date:
      res.append(begin_date)
      begin_date += add_delta
  return res


def convert_string_to_date(date_specifier):
  l = convert_string_to_dates(date_specifier)
  assert len(l) == 1
  return l[0]


def iterate_date(start_date: datetime.date, end_date: datetime.date, reverse=False):
  cur_date = start_date
  if reverse:
    while cur_date > end_date:
      yield cur_date
      cur_date -= datetime.timedelta(days=1)
  else:
    while cur_date < end_date:
      yield cur_date
      cur_date += datetime.timedelta(days=1)


def convert_string_to_datetime(dt_str: str) -> Optional[datetime.datetime]:
  dt = None
  dt_formats = ['%Y%m%d', '%Y%m%dT%H%M%S']
  for dt_format in dt_formats:
    try:
      dt = datetime.datetime.strptime(dt_str, dt_format).replace(tzinfo=pytz.UTC)
      break
    except:
      pass
  assert dt is not None
  return dt


def get_min_max_dt(trading_date: datetime.date
) -> Tuple[datetime.datetime, datetime.datetime]:
  return (
    datetime.datetime.combine(
      trading_date, datetime.time.min.replace(tzinfo=pytz.UTC)),
    datetime.datetime.combine(
      trading_date, datetime.time.max.replace(tzinfo=pytz.UTC))
  )


def get_dt_ranges(start_dt: datetime.datetime,
                  end_dt: datetime.datetime,
                  aggregate: bool=False
) -> List[Tuple[datetime.datetime, datetime.datetime]]:
  if aggregate:
    dt_ranges = [(start_dt, end_dt)]
  else:
    start_date = start_dt.date()
    end_date = end_dt.date()
    if start_date == end_date:
      dt_ranges = [(start_dt, end_dt)]
    else:
      dt_ranges = [(start_dt, get_min_max_dt(start_dt.date())[1])]
      trading_dates = [td for td in iterate_date(start_date + datetime.timedelta(days=1),
                                                 end_date)]
      dt_ranges += [get_min_max_dt(td) for td in trading_dates]
      dt_ranges += [(get_min_max_dt(end_dt.date())[0], end_dt)]
  dt_ranges = [
      (start_dt, end_dt) for start_dt, end_dt in dt_ranges if start_dt != end_dt
  ]
  return dt_ranges
